
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

import matplotlib

# 设置matplotlib的字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体为黑体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题

def plot_costs_and_times(costs, times):
    """
    绘制运输费用和平均等待时间随迭代次数变化的曲线图。

    参数:
    costs -- 运输费用的Y值列表
    times -- 平均等待时间的Y值列表

    返回:
    无
    """
    # 检查输入列表长度是否一致
    if len(costs) != len(times):
        raise ValueError("运输费用和平均等待时间的列表长度必须一致")

    # 生成迭代次数X值
    iterations = range(1, len(costs) + 1)

    # 创建图形和轴对象
    fig, ax = plt.subplots()

    # 绘制运输费用曲线
    ax.plot(iterations, costs, label='运输费用')

    # 绘制平均等待时间曲线
    ax.plot(iterations, times, label='平均等待时间')

    # 设置图表标题和坐标轴标签
    ax.set_title('运输费用和平均等待时间随迭代次数变化')
    ax.set_xlabel('迭代次数')
    ax.set_ylabel('值')

    # 添加图例
    ax.legend()

    # 显示网格
    ax.grid(True)

    # 显示图表
    plt.show()


def generate_city_graph(n, show_img):
    """
    :param show_img: 是否显示图片
    :param n:代表有多少个城市。
    :return: city_ids表示各个城市的编号，customer表示各个城市客户数量，time_matrix记录了各个城市之间的运输时间。
    """
    # 生成城市编号向量
    city_ids = np.arange(0, n) #从0开始

    # 生成随机客户数量
    customer = np.random.randint(3, 20, size=n)

    # 生成一个完全图
    G = nx.complete_graph(n)

    # 生成时间矩阵
    time_matrix = np.full((n, n), -1)
    for i in range(n):
        for j in range(i + 1, n):
            time_matrix[i][j] = time_matrix[j][i] = np.random.randint(3, 20)

    if show_img:
        # 使用 networkx 绘制图 G
        plt.figure(figsize=(5, 5))  # 设置图形的大小

        # 为图 G 设置布局，这里使用 spring_layout 布局
        pos = nx.spring_layout(G)

        # 绘制节点，并在节点上显示城市编号和人口数量
        nx.draw_networkx_nodes(G, pos, node_size=700, node_color='skyblue')
        nx.draw_networkx_labels(G, pos, labels={i: f"{city_ids[i]} ({customer[i]})" for i in range(n)})

        # 绘制边，并在边上显示距离
        edge_labels = {(i, j): time_matrix[i, j] for i in range(n) for j in range(i + 1, n) if
                       time_matrix[i, j] != -1}
        nx.draw_networkx_edges(G, pos, edge_color='gray', width=2)
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color='red')

        plt.axis('off')  # 关闭坐标轴
        plt.show()  # 显示图形

    return city_ids, customer, time_matrix, G


def calculate_fitness(all_path, customer, time_matrix):
    """
    计算每条路径在旅行商问题中的表现，基于平均每个人等待时间和总运输费用。

    参数：
    - all_path: 访问顺序的列表，包含多个路径。
    - customer: 各个城市的人口数量，比如customer[0]表示第一个城市有多少客户，其中0是城市的编号。
    - time_matrix: 城市间的时间矩阵。

    返回：
    - fitness_values: 每个路径在两个目标上的表现，包括平均每个人等待时间和总运输费用。
    """

    fitness_values = []

    for path in all_path:
        total_travel_time = 0  # 该“路径”对应的从头到尾运输总时间，用来计算运费。

        current_time = 0 #当前已经花了多少时间。
        last_city = None #上一个城市的编号。

        cus_and_time=[] #记录各组的人等待了多少时间，其中的每个元素是 [人数，等待时间]，比如[10,33]表示有10个人等待了33单位的时间。

        for city in path[:-1]: #city 是当前所在城市的编号。
            if last_city is None: #表明当前是第一个城市
                # 第一个城市，也就是起点，这里客户平均等待时间为在该城市卸货时间的一半。
                cus_and_time.append([customer[city],customer[city]/2])
                current_time += customer[city] #在该城市卸货花了那么多时间
                last_city = city
            else: #当前不是第一个城市
                current_time += time_matrix[last_city][city] #把路上花的时间计算进来。
                total_travel_time += time_matrix[last_city][city] #更新运输的总时间。

                # 这里客户平均等待时间为在该城市卸货时间的一半加上current_time。
                cus_and_time.append([customer[city], customer[city] / 2 + current_time])
                current_time += customer[city] #在该城市卸货花了那么多时间
                last_city = city

        #还要计算从最后一个城市返回起点的运输时间
        total_travel_time += time_matrix[last_city][0]  # 更新运输的总时间。

        all_time = 0
        all_people = 0

        for item in cus_and_time:
            all_people += item[0]
            all_time += item[1]*item[0]

        average_waiting_time = all_time/all_people
        all_cost = total_travel_time #总运费

        fitness_values.append([average_waiting_time,all_cost])

    return fitness_values

"""
按照如下要求使用蚁群优化算法解决这个双目标的旅行商问题：
采用循环优化的方法，也就是先花n步 优化运输费用，然后花n步 优化客户等待时间 如此循环下去.... 直到达到最大迭代次数。

在优化单个目标的情况下流程如下：
1.使用generate_city_graph函数初始化一个地图(设定城市的数量一共有15个。)。
2.初始化一个模型参数：
蚂蚁数量m：设置为15
信息素矩阵τ：每个元素τ[i][j]记录了城市i和城市j之间的信息素浓度，初始化值从标准正态分布采样得到。
启发式信息矩阵 𝜂: 示运如果是求最佳运输成本对应的路线，那么使用1/time_matrix[i][j]，表输时间小的路径吸引力强。如果是求客户平均等待时间少的路线，使用customer[j]/all_customers，表示客户数量多的城市吸引力更强。
信息素挥发率ρ：用来控制挥发的程度，设置为0.2 。
3.蚂蚁构建路径:蚂蚁在选择下一个城市的时候是以一个概率分布进行选择的,比如从城市i到城市j概率计算使用 τ[i][j]*𝜂[i][j]/all_allowed(𝜂[i][j]),不用考虑什么重要性因子。当蚂蚁经过一条路径之后，该路径上的信息素浓度会增加0.01%。
4.信息素浓度的全局更新：首先，每条路径都按照信息素挥发率ρ=0.2对该路径上的信息素进行衰减，第二步，把所有蚂蚁找到的最优路径上的信息素浓度按照比例p1=0.5%进行额外增加。
5.循环迭代的终止条件：总共最大迭代次数1万次，超过就结束迭代。
6.输出当前找到的最优解
"""

def ant_colony_optimization_cost(n, customer, time_matrix, tau, m, max_iterations, rho=0.2, p1=0.05):
    # 运输费用的优化部分===================================================================================================
    global best_cost
    global best_time
    global best_cost_list
    global best_time_list

    # 初始化启发式信息矩阵 (1/time_matrix[i][j])
    eta = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            eta[i][j] = eta[j][i] = 1 / time_matrix[i][j]  # 对于求运输成本最优的启发式信息

    # 对 eta 归一化
    max_eta = eta.max()  # 找到最大值
    min_eta = eta[eta > 0].min()  # 找到最小的正值，忽略对角线的零值

    eta = (eta - min_eta) / (max_eta - min_eta) + 0.1 # 归一化


    # 迭代开始
    for iteration in range(max_iterations):
        all_paths = []  # 存储所有蚂蚁的路径

        # 蚂蚁构建路径
        for ant in range(m):
            path = [0]  # 该蚂蚁经过的路径记录
            visited = set(path)
            current_city = 0

            while len(path) < n:
                # 计算可行城市的选择概率
                probabilities = []
                for next_city in range(n):
                    if next_city not in visited:

                        prob = tau[current_city][next_city] * eta[current_city][next_city]

                        probabilities.append(prob)
                    else:
                        probabilities.append(0)

                total_prob = sum(probabilities)

                # 归一化概率

                probabilities = [p / total_prob for p in probabilities]

                next_city = np.random.choice(range(n), p=probabilities)

                if tau[current_city][next_city] < 1000:
                    tau[current_city][next_city] *= 1.0001

                path.append(next_city)
                visited.add(next_city)
                current_city = next_city

            path.append(0) #把起点加上去

            all_paths.append(path)

        # 计算所有路径的运输费用
        fitness_values = calculate_fitness(all_paths, customer, time_matrix)

        # 更新最优解
        for i, fitness in enumerate(fitness_values):
            time,cost = fitness[0],fitness[1]  # 运输费用
            if cost < best_cost and time < best_time*1.1:  # 优化运输费用
                best_cost = cost
                best_time = time
                best_path = all_paths[i]

                # 根据最优路径增加信息素
                for index in range(len(best_path) - 1):
                    city1, city2 = best_path[index], best_path[index + 1]
                    tau[city1][city2] += p1 * tau[city1][city2]  # 按照运输费用增加信息素
                    tau[city2][city1] = tau[city1][city2] = min(max(tau[city1][city2], 0.000001), 1000)


        # 全局更新信息素
        for i in range(n):
            for j in range(i + 1, n):
                tau[i][j] *= (1 - rho)  # 衰减信息素
                tau[j][i] = tau[i][j]  # 保证对称

                # 设置下限
                tau[i][j] = max(tau[i][j], 0.000001)
                tau[j][i] = tau[i][j]  # 保证对称

        if best_time < 10000000 and best_cost < 10000000:
            best_cost_list.append(best_cost)
            best_time_list.append(best_time)

    return tau


def ant_colony_optimization_time(n, customer, time_matrix, tau, m, max_iterations, rho=0.2, p1=0.05):
    # 运输费用的优化部分===================================================================================================
    global best_cost
    global best_time
    global best_cost_list
    global best_time_list

    # 初始化启发式信息矩阵
    all_customers = sum(customer)
    eta = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            eta[i][j]  = (customer[j] / all_customers)
            eta[j][i]  = (customer[i] / all_customers)

    # 对 eta 归一化
    max_eta = eta.max()  # 找到最大值
    min_eta = eta[eta > 0].min()  # 找到最小的正值

    eta = (eta - min_eta) / (max_eta - min_eta) + 0.1


    # 迭代开始
    for iteration in range(max_iterations):
        all_paths = []  # 存储所有蚂蚁的路径

        # 蚂蚁构建路径
        for ant in range(m):
            path = [0]  # 该蚂蚁经过的路径记录
            visited = set(path)
            current_city = 0

            while len(path) < n:
                # 计算可行城市的选择概率
                probabilities = []
                for next_city in range(n):
                    if next_city not in visited:

                        prob = tau[current_city][next_city] * eta[current_city][next_city]

                        probabilities.append(prob)
                    else:
                        probabilities.append(0)

                total_prob = sum(probabilities)

                # 归一化概率
                probabilities = [p / total_prob for p in probabilities]

                next_city = np.random.choice(range(n), p=probabilities)

                if tau[current_city][next_city]<1000:
                    tau[current_city][next_city]*=1.0001

                path.append(next_city)
                visited.add(next_city)
                current_city = next_city

            path.append(0) #把起点加上去

            all_paths.append(path)


        fitness_values = calculate_fitness(all_paths, customer, time_matrix)

        # 更新最优解
        for i, fitness in enumerate(fitness_values):
            time,cost = fitness[0],fitness[1]  # 运输费用
            if time < best_time and cost < best_cost*1.1:  # 优化时间
                best_cost = cost
                best_time = time
                best_path = all_paths[i]

                # 根据最优路径增加信息素
                for index in range(len(best_path) - 1):
                    city1, city2 = best_path[index], best_path[index + 1]
                    tau[city1][city2] += p1 * tau[city1][city2]  # 按照运输费用增加信息素
                    tau[city2][city1] = tau[city1][city2] = min(max(tau[city1][city2], 0.000001), 1000)


        # 全局更新信息素
        for i in range(n):
            for j in range(i + 1, n):
                tau[i][j] *= (1 - rho)  # 衰减信息素
                tau[j][i] = tau[i][j]  # 保证对称

                # 设置下限
                tau[i][j] = max(tau[i][j], 0.000001)
                tau[j][i] = tau[i][j]  # 保证对称

        if best_time<10000000 and best_cost <10000000:
            best_cost_list.append(best_cost)
            best_time_list.append(best_time)

    return tau



# 示例运行
n = 15  # 城市数量
m = 15 # 蚂蚁数量

best_cost_list=[] #优化过程当中的最佳运输成本记录。
best_time_list=[] #优化过程当中的最佳等待时间的记录。


# 生成城市图
city_ids, customer, time_matrix, G = generate_city_graph(n, show_img=False)

best_time = 10000000
best_cost = 10000000

# 初始化信息素矩阵
tau = np.random.normal(0, 1, (n, n))  # 信息素浓度矩阵（从标准正态分布采样）

for i in range(n):
    for j in range(i, n):

        tau[j][i] = tau[i][j] = abs(tau[i][j]) # 使 tau[i][j] == tau[j][i]
        if i == j:
            tau[i][j] = 0

train_times = 10000 #总的迭代多少次

done = 0 #已经训练了多少次

while done < train_times:
    # 交替优化 运输时间 和 等待时间
    tau = ant_colony_optimization_cost(n=n,customer=customer,time_matrix=time_matrix,tau=tau,m=m,max_iterations=20)

    tau = ant_colony_optimization_time(n=n, customer=customer, time_matrix=time_matrix, tau=tau, m=m, max_iterations=20)

    done+=40


plot_costs_and_times(costs=best_cost_list,times=best_time_list)

print(f"客户等待时间初始值：{best_time_list[0]},客户等待时间最终值：{best_time_list[-1]}")

print(f"运输费用初始值：{best_cost_list[0]},运输费用最终值：{best_cost_list[-1]}")